# -*- coding: utf-8 -*-
# @Time    : 2020/6/17 上午1:19
# @Author  : caotian
# @FileName: optimizationdata.py
# @Software: PyCharm
import os
import random
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import numpy as np
from PIL import Image
import json
import gzip

def load_data(model='train'):
    datafile='./mnist.json.gz'
    print("loading data from {}".format(datafile))
    data =json.load(gzip.open(datafile))
    train_set,val_set,eval_set=data
    IMG_ROWS=28
    IMG_COLS=28
    if model == "train":
        imgs=train_set[0]
        labels=train_set[1]
    if model == "vaild":
        imgs=val_set[0]
        labels=val_set[1]
    if model == "eval":
        imgs=eval_set[0]
        labels=eval_set[1]
    imgs_length=len(imgs)
    assert len(imgs) == len(labels),"length of imgs {} should be the same as labels {}".format(len(imgs),len(labels))
    index_list=list(range(imgs_length))
    batch_size=100
    def data_generator():
        if model=="train":
            random.shuffle(index_list)
        imgs_list=[]
        labels_list=[]
        for i in index_list:
            img=np.reshape(imgs[i],[1,IMG_ROWS,IMG_COLS]).astype('float32')
            label=np.reshape(labels[i],[1]).astype('int64')
            imgs_list.append(img)
            labels_list.append(label)
            if len(imgs_list) == batch_size:
                yield np.array(imgs_list),np.array(labels_list)
                imgs_list=[]
                labels_list=[]
        if len(imgs_list) > 0:
            yield np.array(imgs_list),np.array(labels_list)
    return data_generator
